window.onload = function () {
    window.scrollTo(0, document.body.scrollHeight);
};

let currentHistoryMessageIdList = [];
let messages = [];

function createMessageElement(message) {
    const div = document.createElement('div');
    div.className = `${message.MsgOwner}MessageContainer messageContainer`;
    div.setAttribute('data-id', message.Id);

    let messageContent;
    if (message.MsgOwner === "user") {
        messageContent = FileDisplay(message.Content);
    }
    else if (message.MsgOwner === "assistant") {
        messageContent = ToLatex(ToMarkDown(message.Content));
    }
    else {
        messageContent = message.Content;
    }

    div.innerHTML = `
        <div class='author ${message.MsgOwner}'>${message.Author}</div>
        <div class='${message.MsgOwner}Message message'>${messageContent}</div>
    `;

    div.querySelectorAll('code').forEach((el) => {
        hljs.highlightElement(el);
    });

    const copyButtons = div.querySelectorAll(".copy-button");
    if (copyButtons.length) {
        copyButtons.forEach(el => {
            el.addEventListener('click', (e) => {
                const code = e.currentTarget.parentNode.nextElementSibling.innerText;
                navigator.clipboard.writeText(code).then(() => {
                    // alert('Code copied to clipboard');
                }).catch(err => {
                    console.error('Failed to copy: ', err);
                });
            });
        });
    }

    let userMessage = div.querySelector(".userMessage.message");
    if (userMessage) {
        addEditMessageFeature(userMessage, message, div);
    } else if (!userMessage) {
        addResendCopyFeature(div, message);
    }

    if (message.MsgOwner === 'tool') {
        const firstLineContent = messageContent.split('\n')[0];

        div.innerHTML = `
            <div class='author ${message.MsgOwner}'>${message.Author}</div>
            <div class='${message.MsgOwner}Message message'>
                <div class='message-preview'>${firstLineContent}</div>
                <div class='message-full'>${messageContent}</div>
                <button class='toggle-button'>展开</button>
            </div>
        `;
        addToolMessageToggleFeature(div);
    }

    if (message.History.length > 0 && userMessage) {
        const lastMessage = messages.find(msg => msg.Id === message.History[message.History.length - 1]);

        if (lastMessage.branchCount > 1) {
            const filteredMessages = messages.filter(msg => msg.History[msg.History.length - 1] === message.History[message.History.length - 1]);

            filteredMessages.sort((a, b) => a.Id - b.Id);

            const rank = filteredMessages.findIndex(item => item.Id === message.Id) + 1;

            addBranchFeature(div, message, rank, lastMessage.branchCount);
        }
    }
    else if (message.History.length === 0) {
        const filteredMessages = messages.filter(msg => msg.History.length === 0);

        if (filteredMessages.length > 1) {
            filteredMessages.sort((a, b) => a.Id - b.Id);

            const rank = filteredMessages.findIndex(item => item.Id === message.Id) + 1;

            addBranchFeature(div, message, rank, filteredMessages.length);
        }
    }
    
    if (/\$.*?\$|\(.*?\)|\{.*?\}/.test(message.Content) && !userMessage) {
        MathJax.typeset([div]);
    }
    
    return div;
}

function addToolMessageToggleFeature(div) {
    const toggleButton = div.querySelector('.toggle-button');
    const messagePreview = div.querySelector('.message-preview');
    const messageFull = div.querySelector('.message-full');

    messageFull.style.display = 'none';
    messagePreview.style.display = 'block';
    toggleButton.textContent = '展开';

    toggleButton.addEventListener('click', () => {
        if (messageFull.style.display === 'none') {
            messageFull.style.display = 'block';
            messagePreview.style.display = 'none';
            toggleButton.textContent = '收起';
        } else {
            messageFull.style.display = 'none';
            messagePreview.style.display = 'block';
            toggleButton.textContent = '展开';
        }
    });
}


function addMessage(message) {
    message.branchCount = 0;

    if (message.History.length !== 0) {
        const lastMessage = messages.find(msg => msg.Id === message.History[message.History.length - 1]);
        if (lastMessage) {
            lastMessage.branchCount += 1;
        }
    }

    messages.push(message);
}

function clearMessage() {
    currentHistoryMessageIdList = [];
    messages = [];
    var allMessages = document.querySelectorAll('div[data-id]');
    allMessages.forEach(function (messageDiv) {
        messageDiv.remove();
    });
}

function refreshDisplayMessages(lastMessageId) {
    let message = messages.find(msg => msg.Id === lastMessageId);

    let historyMessageIdList = message ? JSON.parse(JSON.stringify(message.History)) : [];
    historyMessageIdList.push(lastMessageId);

    currentHistoryMessageIdList.forEach(currentId => {
        if (!historyMessageIdList.includes(currentId)) {
            let message = messages.find(msg => msg.Id === currentId);
            if (message) {
                hideMessage(message);
            }
        }
    });
    
    historyMessageIdList.forEach(id => {
        if (!currentHistoryMessageIdList.includes(id)) {
            let message = messages.find(msg => msg.Id === id);
            if (message) {
                displayMessage(message);
            }
        }
    });

    currentHistoryMessageIdList = [...historyMessageIdList];

    sortMessagesById(messagesContainer);
    window.scrollTo(0, document.body.scrollHeight);
}

function hideMessage(message) {
    var messageDiv = document.querySelector(`div[data-id="${message.Id}"]`);

    if (messageDiv) {
        messageDiv.remove();
    }
}

function displayMessage(message) {
    if (!document.querySelector(`[data-id="${message.Id}"]`)) {
        const messageElement = createMessageElement(message);

        var messagesContainer = document.getElementById('messagesContainer');
        if (messagesContainer) {
            messagesContainer.appendChild(messageElement);
        }
    }
}

function updateMessage(message) {
    var messageDiv = document.querySelector(`div[data-id="${message.Id}"]`);
    if (messageDiv) {
        var contentDiv = messageDiv.querySelector(`.${message.MsgOwner}Message.message`);
        
        if (contentDiv) {
            const lastMessage = currentHistoryMessageIdList.length > 0
                ? messages.find(msg => msg.Id === currentHistoryMessageIdList[currentHistoryMessageIdList.length - 1])
                : null;

            if (lastMessage == null || (message.Id === lastMessage.Id && message.Content.length > lastMessage.Content.length)) {
                const newMessageElement = createMessageElement(message);

                messageDiv.replaceWith(newMessageElement);

                lastMessage.Content = message.Content;

                scrollToBottomIfNearBottom();
            }
        }
    }
}

function sortMessagesById(container) {
    var messages = Array.from(container.children);
    messages.sort(function (a, b) {
        return parseInt(a.getAttribute('data-id')) - parseInt(b.getAttribute('data-id'));
    });
    messages.forEach(function (msg) {
        container.appendChild(msg);
    });
}

function scrollToBottomIfNearBottom() {
    const currentScrollPosition = window.scrollY;
    const totalPageHeight = document.body.scrollHeight;
    const windowHeight = window.innerHeight;
    const bottomThreshold =80;
    const nearBottom = (totalPageHeight - (currentScrollPosition + windowHeight)) < bottomThreshold;

    if (nearBottom) {
        window.scrollTo(0, totalPageHeight);
    }
}

function addEditMessageFeature(userMessage, message, div) {
    let editIcon = document.createElement('span');
    editIcon.className = 'edit-icon';
    editIcon.innerHTML = '&#9998;';

    userMessage.style.position = 'relative';
    userMessage.appendChild(editIcon);

    editIcon.addEventListener('click', () => {
        if (!div.classList.contains('editing')) {
            div.classList.add('editing');
            let tempUserMessage = userMessage.cloneNode(true);
            let tempEditIcon = tempUserMessage.querySelector('.edit-icon');
            if (tempEditIcon) {
                tempEditIcon.remove();
            }
            const originalContent = tempUserMessage.textContent.trim();

            const editInput = document.createElement('textarea');
            editInput.className = 'edit-input';
            editInput.value = originalContent;
            editInput.addEventListener('input', () => {
                editInput.style.height = 'auto'
                editInput.style.height = (editInput.scrollHeight) + 'px'
            });

            const editButtons = document.createElement('div');
            editButtons.className = 'edit-buttons';

            const cancelButton = document.createElement('button');
            cancelButton.textContent = '取消';
            cancelButton.className = 'cancel-button';

            const saveButton = document.createElement('button');
            saveButton.textContent = '覆盖';
            saveButton.className = 'save-message-button';

            const sendButton = document.createElement('button');
            sendButton.textContent = '发送';
            sendButton.className = 'send-button';

            editButtons.appendChild(cancelButton);
            editButtons.appendChild(saveButton);
            editButtons.appendChild(sendButton);

            userMessage.innerHTML = '';
            userMessage.appendChild(editInput);
            userMessage.appendChild(editButtons);
            editInput.focus();


            cancelButton.addEventListener('click', () => {
                userMessage.innerHTML = originalContent;
                addEditMessageFeature(userMessage, message, div);
                div.classList.remove('editing');
            });

            saveButton.addEventListener('click', () => {
                const newContent = editInput.value.trim();
                message.Content = newContent;
                updateMessage(message);
                userMessage.innerHTML = newContent;
                addEditMessageFeature(userMessage, message, div);
                div.classList.remove('editing');

                overrideMessage(message.Id, newContent)
            });

            sendButton.addEventListener('click', () => {
                const newContent = editInput.value.trim();

                hideMessage(message)

                resendMessage(message, newContent);
            });
        }
    });
}

function addBranchFeature(div, message, number1, number2) {
    
    const actionIconsContainer = document.createElement('div');
    actionIconsContainer.className = 'action-icons-container';

    let leftBranch = document.createElement('span');
    leftBranch.className = 'left-branch';
    leftBranch.innerHTML = '&lt;';
    actionIconsContainer.appendChild(leftBranch);

    let numberDisplay = document.createElement('span');
    numberDisplay.className = 'number-display';
    numberDisplay.innerHTML = `${number1}/${number2}`;
    actionIconsContainer.appendChild(numberDisplay);

    let rightBranch = document.createElement('span');
    rightBranch.className = 'right-branch';
    rightBranch.innerHTML = '&gt;';
    actionIconsContainer.appendChild(rightBranch);

    div.appendChild(actionIconsContainer);

    leftBranch.addEventListener('click', () => {
        switchBranch(message.History.length > 0 ? message.History[message.History.length - 1] : 0, message.Id, -1);
    });

    rightBranch.addEventListener('click', () => {
        switchBranch(message.History.length > 0 ? message.History[message.History.length - 1] : 0, message.Id, 1);
    });
}

function switchBranch(lastId, curId, leftOrRight) {
    const filteredMessages = lastId === 0 ?
        messages.filter(message =>
            message.History.length === 0
        )
        : messages.filter(message =>
            message.History.length > 0 && message.History[message.History.length - 1] === lastId
        );

    filteredMessages.sort((a, b) => a.Id - b.Id);

    let index = filteredMessages.findIndex(message => message.Id === curId);
    if (index === -1) return null;

    let nextIndex = (filteredMessages.length + index + leftOrRight) % filteredMessages.length;

    let prefixHistory = filteredMessages[nextIndex].History;

    let longestHistory = JSON.parse(JSON.stringify(prefixHistory));
    longestHistory.push(filteredMessages[nextIndex].Id);


    for (let message of messages) {
        let history = message.History;
        if (history.length >= longestHistory.length &&
            history.slice(0, longestHistory.length).toString() === longestHistory.toString()) {
            longestHistory = JSON.parse(JSON.stringify(history));
            longestHistory.push(message.Id);
        }
    }

    const resendPayload = {
        messageType: "switchBranch",
        params: {
            messageId: longestHistory[longestHistory.length - 1]
        }
    };

    window.chrome.webview.postMessage(JSON.stringify(resendPayload));

    return longestHistory[longestHistory.length - 1];
}


function addResendCopyFeature(div, message) {
    const actionIconsContainer = document.createElement('div');
    actionIconsContainer.className = 'action-icons-container';

    let copyIcon = document.createElement('span');
    copyIcon.className = 'copy-icon';
    copyIcon.innerHTML = '&#128203;';
    actionIconsContainer.appendChild(copyIcon);

    let resendIcon = document.createElement('span');
    resendIcon.className = 'resend-icon';
    resendIcon.innerHTML = '&#8635;';
    actionIconsContainer.appendChild(resendIcon);

    div.appendChild(actionIconsContainer);

    resendIcon.addEventListener('click', () => {
        resendMessage(message, null);
    });

    copyIcon.addEventListener('click', () => {
        const messageContent = message.Content;
        navigator.clipboard.writeText(messageContent).then(() => {
            console.log('Message copied to clipboard');
        }).catch(err => {
            console.error('Failed to copy: ', err);
        });
    });
}





function resendMessage(message, content) {
    const resendPayload = {
        messageType: "resendMessage",
        params: {
            messageId: message.Id,
            content: content
        }
    };

    for (let id of currentHistoryMessageIdList) {
        if (id >= message.Id) {
            hideMessage(messages.find(msg => msg.Id === id));
        }
    }

    window.chrome.webview.postMessage(JSON.stringify(resendPayload));
}

function overrideMessage(messageId, content) {
    if (content === "") {
        return;
    }

    const resendPayload = {
        messageType: "overrideMessage",
        params: {
            messageId: messageId,
            content: content
        }
    };
    
    window.chrome.webview.postMessage(JSON.stringify(resendPayload));
}
